"""
Case Type   : DiskANN向量检索
Case Name   : diskann索引重建过程进行DML、vacuum、analyze等操作，功能正常
Create At   : 2025-09-19
Owner       : taoying
Description :
    1.创建表格并插入数据、构建索引
    2.循环进行vacuum/update操作
    3.循环进行索引重建
Expect      :
    1.成功
    2.成功
    3.成功
History     :
"""
import os
import unittest

from yat.test import Node
from yat.test import macro

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Logger import Logger

logger = Logger()


class DiskANNSearch(unittest.TestCase):
    def setUp(self):
        logger.info('----------------this is setup-----------------------')
        logger.info(
            '---Opengauss_Function_DiskANN_Index_Case0036开始执行-----')
        self.userNode = Node('dbuser')
        self.DB_ENV_PATH = macro.DB_ENV_PATH
        self.comsh = CommonSH('dbuser')
        self.com = Common()
        self.comsql = CommonSQL('dbuser')
        self.table_name = 't_func_diskann_index0036'
        self.fun_insert = 'f_func_diskann_index0036'

    def test_diskann(self):
        step_txt = '----step1:创建表格和diskann索引; expect:成功----'
        sql = f'''
            drop table if exists {self.table_name};
            create table {self.table_name} (id bigserial primary key, val1 vector(5));
            alter table {self.table_name} set (parallel_workers=10);
            drop index if exists idx_t_func_diskann_index0036;
            create index idx_t_func_diskann_index0036 on {self.table_name} using diskann(val1 vector_l2_ops) with(index_size=16);
            '''
        sql_cmd = self.comsh.execut_db_sql(sql)
        logger.info(sql_cmd)

        step_txt = '----step2:创建自定义插入函数; expect:成功----'
        logger.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        logger.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step3:通过自定义插入函数插入数据; expect:成功----'
        insert_sql = f"select {self.fun_insert}(5, 100, '{self.table_name}', 'val1', 'vector');"
        result = self.comsh.execut_db_sql(insert_sql)
        logger.info(result)
        self.assertIn('SUCCESS', result, '执行失败: 插入数据失败')

        step_txt = '----step4:循环进行vacuum/update操作; expect:成功----'
        excute_cmd1 = '''source {DB_ENV_PATH};
        for i in $(seq 1 10);
        do gsql -d test_db \\
        -p {db_port} \\
        -U {ssh_user} \\
        -c "delete from t_func_diskann_index0036 where id<10; vacuum full t_func_diskann_index0036; update t_func_diskann_index0036 set val1=val1-'[1,1,1,1,1]' where id<10;"
        done;
        '''.format(
            DB_ENV_PATH=self.DB_ENV_PATH,
            db_port=self.userNode.db_port,
            ssh_user=self.userNode.ssh_user
        )
        logger.info(excute_cmd1)
        msg1 = self.userNode.sh(excute_cmd1).result()
        logger.info(msg1)
        self.assertNotIn('Error', msg1)

        step_txt = '----step5:循环进行索引重建; expect:成功----'
        excute_cmd2 = f'''source {self.DB_ENV_PATH};
        for i in $(seq 1 10);
        do gsql -d test_db \\
        -p {self.userNode.db_port} \\
        -U {self.userNode.ssh_user} \\
        -c "reindex index idx_t_func_diskann_index0036;"
        done;
        '''
        logger.info(excute_cmd2)
        msg2 = self.userNode.sh(excute_cmd2).result()
        logger.info(msg2)
        self.assertNotIn('Error', msg2)

    def tearDown(self):
        logger.info('----------------this is tearDown-----------------------')
        sql_cmd = self.comsh.execut_db_sql('''drop table t_func_diskann_index0036;''')
        logger.info(sql_cmd)
        logger.info(
            '----Opengauss_Function_DiskANN_Index_Case0036执行完成-----')
